// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <utility>
#include <pollux/common/file/file_systems.h>
#include <pollux/tool/trace/operator_replayer_base.h>

#include <melon/executors/io_thread_pool_executor.h>
#include <pollux/exec/task_trace_reader.h>

TURBO_DECLARE_FLAG(std::string, root_dir);
TURBO_DECLARE_FLAG(bool, summary);
TURBO_DECLARE_FLAG(bool, short_summary);
TURBO_DECLARE_FLAG(std::string, query_id);
TURBO_DECLARE_FLAG(std::string, task_id);
TURBO_DECLARE_FLAG(std::string, node_id);
TURBO_DECLARE_FLAG(int32_t, driver_id);
TURBO_DECLARE_FLAG(std::string, driver_ids);
TURBO_DECLARE_FLAG(std::string, table_writer_output_dir);
TURBO_DECLARE_FLAG(double, hive_connector_executor_hw_multiplier);
TURBO_DECLARE_FLAG(int32_t, shuffle_serialization_format);
DECLARE_uint64(query_memory_capacity_mb);
TURBO_DECLARE_FLAG(double, driver_cpu_executor_hw_multiplier);
TURBO_DECLARE_FLAG(std::string, memory_arbitrator_type);
TURBO_DECLARE_FLAG(bool, copy_results);
TURBO_DECLARE_FLAG(std::string, function_prefix);

namespace kumo::pollux::tool::trace {

/// The trace replay runner. It is configured through a set of gflags passed
/// from replayer tool command line.
class TraceReplayRunner {
 public:
  TraceReplayRunner();
  virtual ~TraceReplayRunner() = default;

  /// Initializes the trace replay runner by setting the pollux runtime
  /// environment for the trace replay. It is invoked before run().
  virtual void init();

  /// Runs the trace replay with a set of gflags passed from replayer tool.
  virtual void run();

 protected:
  std::unique_ptr<tool::trace::OperatorReplayerBase> createReplayer() const;

  const std::unique_ptr<melon::CPUThreadPoolExecutor> cpuExecutor_;
  const std::unique_ptr<melon::IOThreadPoolExecutor> ioExecutor_;
  std::shared_ptr<filesystems::FileSystem> fs_;
  std::unique_ptr<exec::trace::TaskTraceMetadataReader>
      taskTraceMetadataReader_;
  core::PlanNodePtr tracePlanNode_;
};

} // namespace kumo::pollux::tool::trace
